import re
import cv2
import numpy as np
import os
import shutil
import sys
from pdf2image import convert_from_path
from paddleocr import PaddleOCR
from PIL import Image, ImageDraw, ImageFont
from fpdf import FPDF
import argparse
import logging
from logging.handlers import RotatingFileHandler
import time

# -----------------------------
# 日志系统（保持不变）
# -----------------------------
def setup_logging():
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.DEBUG)
    log_dir = "id_desensitize_logs"
    os.makedirs(log_dir, exist_ok=True)
    log_file = os.path.join(log_dir, "desensitize.log")
    formatter = logging.Formatter(
        "%(asctime)s - %(levelname)s - %(module)s:%(lineno)d - %(message)s"
    )
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(formatter)
    file_handler = RotatingFileHandler(
        log_file, maxBytes=5 * 1024 * 1024, backupCount=5, encoding='utf-8'
    )
    file_handler.setLevel(logging.DEBUG)
    file_handler.setFormatter(formatter)
    logger.addHandler(console_handler)
    logger.addHandler(file_handler)
    return logger

logger = setup_logging()

# -----------------------------
# 配置类（保持不变）
# -----------------------------
class Config:
    POPPLER_PATH = r"D:\Workspace\desensitization\poppler\bin"
    PDF_DPI = 600
    USE_LOSSLESS_FORMAT = True
    TEMP_DIR = "id_desensitize_temp"
    DEBUG_IMG_DIR = os.path.join(TEMP_DIR, "debug")
    ID_PATTERN = re.compile(r"\d{17}[\dXx]|\d{15}|\d{6,8}")
    TITLE_KEYWORDS = {"正高级", "高级", "中级", "初级", "教授", "副教授"}
    NUMBER_PATTERN = re.compile(r"\d{2,}")
    OCR_ARGS = {"use_angle_cls": True, "lang": "ch"}
    FONT_PATH = r"C:\Windows\Fonts\msyh.ttc"
    FONT_SCALE = 0.8
    CHAR_SPACING = 1.3
    SEGMENT_SPACING = 5
    MIN_SEGMENT_GAP = 3
    DARK_THRESHOLD = 50
    OCCLUSION_WIDTH_MIN = 80
    OCCLUSION_HEIGHT_MIN = 20
    OCCLUSION_ASPECT_RATIO = (5, 15)

# -----------------------------
# 工具函数（修改了文件处理部分）
# -----------------------------
def init_environment():
    start_time = time.time()
    logger.debug("开始初始化环境目录")
    os.makedirs(Config.TEMP_DIR, exist_ok=True)
    os.makedirs(Config.DEBUG_IMG_DIR, exist_ok=True)
    file_count = 0
    for f in os.listdir(Config.DEBUG_IMG_DIR):
        try:
            os.remove(os.path.join(Config.DEBUG_IMG_DIR, f))
            file_count += 1
        except Exception as e:
            logger.warning(f"清理旧调试文件失败: {e}")
    logger.debug(f"环境初始化完成，清理了{file_count}个旧调试文件，耗时{time.time()-start_time:.4f}秒")

def save_debug_img(img, name):
    start_time = time.time()
    path = os.path.join(Config.DEBUG_IMG_DIR, name)
    try:
        if Config.USE_LOSSLESS_FORMAT:
            cv2.imencode('.png', img)[1].tofile(path)
        else:
            cv2.imencode('.jpg', img)[1].tofile(path)
        img_height, img_width = img.shape[:2]
        logger.info(f"调试图片已保存: {path} (尺寸: {img_width}x{img_height}px)")
        logger.debug(f"保存调试图片耗时{time.time()-start_time:.4f}秒")
    except Exception as e:
        logger.error(f"保存调试图片失败: {str(e)}", exc_info=True)

def detect_dark_occlusions(image):
    start_time = time.time()
    logger.debug("开始检测黑色遮挡区域")
    try:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        _, dark_mask = cv2.threshold(gray, Config.DARK_THRESHOLD, 255, cv2.THRESH_BINARY_INV)
        contours, _ = cv2.findContours(dark_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        occlusions = []
        for cnt in contours:
            x, y, w, h = cv2.boundingRect(cnt)
            aspect_ratio = w / h
            if (w >= Config.OCCLUSION_WIDTH_MIN and
                    h >= Config.OCCLUSION_HEIGHT_MIN and
                    Config.OCCLUSION_ASPECT_RATIO[0] <= aspect_ratio <= Config.OCCLUSION_ASPECT_RATIO[1]):
                occlusions.append((x, y, x + w, y + h))
                logger.debug(f"检测到黑色遮挡区域: 位置({x}, {y}), 大小({w}x{h}), 宽高比{aspect_ratio:.2f}")
        logger.debug(f"黑色遮挡区域检测完成，共发现{len(occlusions)}个区域，耗时{time.time()-start_time:.4f}秒")
        return occlusions
    except Exception as e:
        logger.error(f"黑色遮挡区域检测失败: {str(e)}", exc_info=True)
        return []

def get_ocr_processed_image(original_img):
    start_time = time.time()
    logger.debug("开始OCR预处理")
    try:
        ocr_img = original_img.copy()
        ocr_img = cv2.GaussianBlur(ocr_img, (3, 3), 0)
        logger.debug(f"OCR预处理完成，耗时{time.time()-start_time:.4f}秒")
        return ocr_img
    except Exception as e:
        logger.error(f"OCR预处理失败: {str(e)}", exc_info=True)
        return original_img.copy()

def split_compound_text(text):
    start_time = time.time()
    logger.debug(f"开始拆分混合文本: {text}")
    segments = []
    current_pos = 0
    text_len = len(text)
    try:
        while current_pos < text_len:
            id_match = Config.ID_PATTERN.search(text, current_pos)
            title_match = None
            for title in Config.TITLE_KEYWORDS:
                title_pos = text.find(title, current_pos)
                if title_pos != -1:
                    title_match = (title_pos, title_pos + len(title), title)
                    break
            num_match = Config.NUMBER_PATTERN.search(text, current_pos)
            matches = []
            if id_match:
                matches.append(('id', id_match.start(), id_match.end()))
            if title_match:
                matches.append(('title', title_match[0], title_match[1]))
            if num_match:
                matches.append(('number', num_match.start(), num_match.end()))
            if not matches:
                segments.append({'type': 'normal', 'content': text[current_pos:]})
                break
            matches.sort(key=lambda x: x[1])
            best_type, best_start, best_end = matches[0]
            if best_start > current_pos:
                segments.append({'type': 'normal', 'content': text[current_pos:best_start]})
            if best_type == 'id':
                segments.append({'type': 'id', 'content': text[best_start:best_end]})
            elif best_type == 'title':
                segments.append({'type': 'title', 'content': text[best_start:best_end]})
            else:
                segments.append({'type': 'number', 'content': text[best_start:best_end]})
            current_pos = best_end
        segment_types = [seg['type'] for seg in segments]
        logger.debug(f"文本拆分完成，得到{len(segments)}个片段，类型: {segment_types}，耗时{time.time()-start_time:.4f}秒")
        return segments
    except Exception as e:
        logger.error(f"文本拆分失败: {str(e)}", exc_info=True)
        return [{'type': 'normal', 'content': text}]

def calculate_segment_positions(segments, full_bbox, font):
    start_time = time.time()
    logger.debug(f"开始计算{len(segments)}个文本片段的位置")
    x1, y1, x2, y2 = full_bbox
    text_height = y2 - y1
    total_width = x2 - x1
    try:
        segment_widths = []
        for seg in segments:
            bbox = ImageDraw.Draw(Image.new('RGB', (1, 1))).textbbox((0, 0), seg['content'], font=font)
            seg_width = bbox[2] - bbox[0]
            segment_widths.append(seg_width)
        total_segment_width = sum(segment_widths)
        total_gap = Config.SEGMENT_SPACING * (len(segments) - 1)
        total_needed_width = total_segment_width + total_gap
        scale = 1.0
        if total_needed_width > total_width:
            scale = total_width / total_needed_width
            logger.info(f"文本总宽度({total_needed_width:.1f}px)超过文本框宽度({total_width:.1f}px)，应用缩放比例: {scale:.2f}")
        current_x = x1
        positions = []
        for i, (seg, width) in enumerate(zip(segments, segment_widths)):
            scaled_width = width * scale
            positions.append({
                **seg,
                'bbox': (current_x, y1, current_x + scaled_width, y2),
                'width': scaled_width
            })
            current_x += scaled_width + (Config.SEGMENT_SPACING * scale if i < len(segments)-1 else 0)
        logger.debug(f"文本片段位置计算完成，缩放比例: {scale:.2f}，耗时{time.time()-start_time:.4f}秒")
        return positions
    except Exception as e:
        logger.error(f"文本片段位置计算失败: {str(e)}", exc_info=True)
        return segments

def detect_id_numbers(original_img, page_num, ocr):
    start_time = time.time()
    logger.info(f"开始检测第{page_num}页的身份证号")
    save_debug_img(original_img, f"page_{page_num}_original.png")
    ocr_img = get_ocr_processed_image(original_img)
    save_debug_img(ocr_img, f"page_{page_num}_ocr_input.png")
    id_list = []
    try:
        result = ocr.ocr(ocr_img, cls=True)
        if not result or not result[0]:
            logger.warning(f"第{page_num}页未识别到文本")
            return id_list
        logger.debug(f"第{page_num}页共识别到{len(result[0])}行文本")
        for line_idx, line in enumerate(result[0]):
            text, confidence = line[1]
            if confidence < 0.3:
                logger.debug(f"第{page_num}页第{line_idx+1}行文本置信度低({confidence:.2f})，跳过: {text}")
                continue
            segments = split_compound_text(text)
            has_id = any(seg['type'] == 'id' for seg in segments)
            if not has_id:
                logger.debug(f"第{page_num}页第{line_idx+1}行文本不包含身份证号，跳过: {text}")
                continue
            coords = line[0]
            x1, y1 = min(p[0] for p in coords), min(p[1] for p in coords)
            x2, y2 = max(p[0] for p in coords), max(p[1] for p in coords)
            font_size = int((y2 - y1) * Config.FONT_SCALE)
            id_list.append({
                "type": "compound_text",
                "full_text": text,
                "full_bbox": (x1, y1, x2, y2),
                "segments": segments,
                "font_size": font_size,
                "confidence": confidence
            })
            logger.info(f"第{page_num}页第{line_idx+1}行识别到含身份证号的文本: {text}（置信度: {confidence:.2f}）")
        occlusions = detect_dark_occlusions(ocr_img)
        for box in occlusions:
            x1, y1, x2, y2 = box
            font_size = int((y2 - y1) * Config.FONT_SCALE)
            star_count = int((x2 - x1) / (font_size * 0.55))
            star_count = max(15, min(18, star_count))
            id_list.append({
                "type": "occlusion",
                "box": (x1, y1, x2, y2),
                "font_size": font_size,
                "star_count": star_count
            })
        logger.info(f"第{page_num}页身份证号检测完成，共发现{len(id_list)}个需要处理的区域，耗时{time.time()-start_time:.2f}秒")
        return id_list
    except Exception as e:
        logger.error(f"第{page_num}页OCR识别出错: {str(e)}", exc_info=True)
        return []

def desensitize_image(original_img, id_info_list):
    start_time = time.time()
    logger.debug(f"开始脱敏处理，共{len(id_info_list)}个区域")
    try:
        pil_img = Image.fromarray(cv2.cvtColor(original_img, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(pil_img)
        for info_idx, info in enumerate(id_info_list):
            if info["type"] == "compound_text":
                full_bbox = info["full_bbox"]
                segments = info["segments"]
                font_size = info["font_size"]
                x1, y1, x2, y2 = full_bbox
                logger.debug(f"处理第{info_idx+1}个混合文本区域: {info['full_text']}，字体大小: {font_size}px")
                bg_x = max(0, min(original_img.shape[1]-1, int(x1) - 10))
                bg_y = max(0, min(original_img.shape[0]-1, int(y1) - 10))
                bg_color = pil_img.getpixel((bg_x, bg_y))
                draw.rectangle([(x1, y1), (x2, y2)], fill=bg_color)
                try:
                    font = ImageFont.truetype(Config.FONT_PATH, font_size, encoding="utf-8")
                    logger.debug(f"成功加载字体: {Config.FONT_PATH}，大小: {font_size}px")
                except (IOError, OSError) as e:
                    logger.warning(f"无法加载指定字体，使用默认字体: {str(e)}")
                    font = ImageFont.load_default()
                segment_positions = calculate_segment_positions(segments, full_bbox, font)
                for seg in segment_positions:
                    if seg["type"] == "id":
                        id_text = seg["content"]
                        if len(id_text) >= 18:
                            display_text = id_text[:6] + "********" + id_text[-4:]
                        elif len(id_text) >= 15:
                            display_text = id_text[:6] + "*****" + id_text[-4:]
                        elif len(id_text) >= 11:
                            display_text = id_text[:3] + "*****" + id_text[-3:]
                        else:
                            display_text = id_text[:3] + "*****"
                        logger.debug(f"身份证号脱敏: {id_text} -> {display_text}")
                    else:
                        display_text = seg["content"]
                    draw.text(
                        (seg["bbox"][0], seg["bbox"][1]),
                        display_text,
                        font=font,
                        fill=(0, 0, 0),
                        stroke_width=0
                    )
            elif info["type"] == "occlusion":
                x1, y1, x2, y2 = info["box"]
                font_size = info["font_size"]
                logger.debug(f"处理第{info_idx+1}个黑色遮挡区域，位置: ({x1}, {y1})-({x2}, {y2})，字体大小: {font_size}px")
                bg_x = max(0, min(original_img.shape[1]-1, int(x1) - 10))
                bg_y = max(0, min(original_img.shape[0]-1, int(y1) - 10))
                bg_color = pil_img.getpixel((bg_x, bg_y))
                draw.rectangle([(x1, y1), (x2, y2)], fill=bg_color)
                desensitized_text = "*" * info["star_count"]
                try:
                    font = ImageFont.truetype(Config.FONT_PATH, font_size)
                except:
                    font = ImageFont.load_default()
                bbox = draw.textbbox((0, 0), desensitized_text, font=font)
                text_width = bbox[2] - bbox[0]
                text_height = bbox[3] - bbox[1]
                text_x = round(x1 + (x2 - x1 - text_width) / 2, 1)
                text_y = round(y1 + (y2 - y1 - text_height) / 2, 1)
                draw.text((text_x, text_y), desensitized_text, font=font, fill=(0, 0, 0))
        result_img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
        logger.debug(f"脱敏处理完成，耗时{time.time()-start_time:.2f}秒")
        return result_img
    except Exception as e:
        logger.error(f"脱敏处理失败: {str(e)}", exc_info=True)
        return original_img

def pdf_to_images(pdf_path):
    start_time = time.time()
    logger.info(f"开始将PDF转换为图片: {os.path.basename(pdf_path)}")
    try:
        # 使用cv2.imdecode读取中文路径
        with open(pdf_path, 'rb') as f:
            pdf_bytes = f.read()
        
        images = convert_from_path(
            pdf_bytes if isinstance(pdf_path, bytes) else pdf_path,
            dpi=Config.PDF_DPI,
            poppler_path=Config.POPPLER_PATH,
            fmt='png' if Config.USE_LOSSLESS_FORMAT else 'jpeg',
            use_cropbox=True,
            thread_count=os.cpu_count()
        )
        result = [np.array(img) for img in images]
        logger.info(f"PDF转换完成，共{len(result)}页，耗时{time.time()-start_time:.2f}秒")
        return result
    except Exception as e:
        logger.error(f"PDF转换失败: {str(e)}", exc_info=True)
        return None

# -----------------------------
# 处理单张图片（修改了文件读取方式）
# -----------------------------
def process_single_image(image_path, output_path, ocr):
    """处理单张图片，保持原格式输出"""
    start_time = time.time()
    logger.info(f"===== 开始处理图片文件: {os.path.basename(image_path)} =====")

    try:
        # 使用cv2.imdecode读取中文路径
        with open(image_path, 'rb') as f:
            img_bytes = np.frombuffer(f.read(), np.uint8)
        original_img = cv2.imdecode(img_bytes, cv2.IMREAD_COLOR)
        
        if original_img is None:
            logger.error(f"无法读取图片: {image_path}")
            return False

        id_info_list = detect_id_numbers(original_img, 1, ocr)
        total_desensitized = len(id_info_list)

        if id_info_list:
            processed_img = desensitize_image(original_img, id_info_list)
            save_debug_img(processed_img, f"image_desensitized.png")
        else:
            logger.info("未找到需要脱敏的内容")
            processed_img = original_img

        # 保持原格式输出
        pil_img = Image.fromarray(cv2.cvtColor(processed_img, cv2.COLOR_BGR2RGB))
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True)
        pil_img.save(output_path, quality=95, optimize=True)
        logger.info(f"图片处理完成，结果保存至: {output_path}")

        total_time = time.time() - start_time
        logger.info(f"===== 图片处理完成！共脱敏 {total_desensitized} 个区域，总耗时{total_time:.2f}秒 =====")
        return True

    except Exception as e:
        logger.error(f"处理图片失败: {str(e)}", exc_info=True)
        return False

# -----------------------------
# 处理单个文件（通用入口）
# -----------------------------
def process_file(input_path, output_path):
    """统一处理入口：支持 PDF 和 图片"""
    _, ext = os.path.splitext(input_path.lower())
    ocr = PaddleOCR(**Config.OCR_ARGS)  # 共享 OCR 实例

    if ext == '.pdf':
        return process_single_pdf(input_path, output_path, ocr)
    elif ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.gif']:
        return process_single_image(input_path, output_path, ocr)
    else:
        logger.error(f"不支持的文件格式: {ext}")
        return False

# -----------------------------
# 处理单个PDF（复用 OCR 实例）
# -----------------------------
def process_single_pdf(input_path, output_path, ocr):
    """处理单个PDF文件（复用OCR实例）"""
    total_start_time = time.time()
    logger.info(f"===== 开始处理PDF文件: {os.path.basename(input_path)} =====")
    init_environment()
    images = pdf_to_images(input_path)
    if not images:
        logger.error(f"无法处理文件: {input_path}")
        return False
    processed_images = []
    total_desensitized = 0
    for page_num, original_img in enumerate(images, 1):
        page_start_time = time.time()
        logger.info(f"开始处理第{page_num}页")
        id_info_list = detect_id_numbers(original_img, page_num, ocr)
        if id_info_list:
            total_desensitized += len(id_info_list)
            logger.info(f"第{page_num}页找到 {len(id_info_list)} 个需要脱敏的区域")
            processed_img = desensitize_image(original_img, id_info_list)
            processed_images.append(processed_img)
            save_debug_img(processed_img, f"page_{page_num}_desensitized.png")
        else:
            logger.info(f"第{page_num}页未找到需要脱敏的内容")
            processed_images.append(original_img)
        logger.info(f"第{page_num}页处理完成，耗时{time.time()-page_start_time:.2f}秒")

    # 生成最终PDF
    pdf_start_time = time.time()
    logger.info(f"开始生成最终PDF: {os.path.basename(output_path)}")
    try:
        pdf = FPDF(unit="mm", format="A4")
        for img in processed_images:
            pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
            temp_path = os.path.join(Config.TEMP_DIR, "temp_page.png")
            pil_img.save(temp_path, format="PNG", quality=98)
            pdf.add_page()
            img_width, img_height = pil_img.size
            pdf_width, pdf_height = 210, 297
            scale = min(pdf_width / (img_width / 25.4), pdf_height / (img_height / 25.4))
            pdf.image(temp_path, x=0, y=0, w=img_width * scale / 25.4, h=img_height * scale / 25.4)
        os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True)
        pdf.output(output_path)
        logger.info(f"PDF生成完成，耗时{time.time()-pdf_start_time:.2f}秒")
    except Exception as e:
        logger.error(f"PDF生成失败: {str(e)}", exc_info=True)
        return False

    try:
        shutil.rmtree(Config.TEMP_DIR, ignore_errors=True)
        logger.debug("临时文件已清理")
    except Exception as e:
        logger.warning(f"清理临时文件失败: {e}")

    total_time = time.time() - total_start_time
    logger.info(f"===== 文件处理完成！共脱敏 {total_desensitized} 个区域，总耗时{total_time:.2f}秒，结果保存至: {output_path} =====")
    return True

# -----------------------------
# 支持文件夹或单文件
# -----------------------------
def process_folder(input_folder, output_folder):
    """处理文件夹中所有支持的文件"""
    start_time = time.time()
    logger.info(f"开始处理文件夹: {input_folder}，输出至: {output_folder}")
    supported_files = []
    for f in os.listdir(input_folder):
        try:
            if f.lower().endswith((".pdf", ".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".gif")):
                supported_files.append(f)
        except:
            continue
    
    total_files = len(supported_files)
    logger.info(f"在输入文件夹中发现{total_files}个支持的文件")
    processed_count = skipped_count = failed_count = 0

    ocr = PaddleOCR(**Config.OCR_ARGS)  # 共享 OCR 实例

    for file_idx, filename in enumerate(supported_files, 1):
        logger.info(f"\n处理文件 {file_idx}/{total_files}: {filename}")
        input_path = os.path.join(input_folder, filename)
        output_path = os.path.join(output_folder, filename)
        if os.path.exists(output_path):
            if not ask_to_override(output_path):
                logger.info(f"跳过文件: {filename}")
                skipped_count += 1
                continue
        if process_file(input_path, output_path):
            processed_count += 1
        else:
            failed_count += 1

    total_time = time.time() - start_time
    logger.info("\n===== 所有文件处理统计 =====")
    logger.info(f"总文件数: {total_files}")
    logger.info(f"成功处理: {processed_count}")
    logger.info(f"已跳过: {skipped_count}")
    logger.info(f"处理失败: {failed_count}")
    logger.info(f"总耗时: {total_time:.2f}秒")
    logger.info("===========================")

def ask_to_override(file_path):
    try:
        response = input(f"文件 '{os.path.basename(file_path)}' 已存在。是否覆盖? (y/n): ").strip().lower()
        logger.debug(f"用户对文件'{os.path.basename(file_path)}'的覆盖选择: {response}")
        if response in ['y', 'yes']:
            return True
        elif response in ['n', 'no']:
            return False
        else:
            print("请输入 'y' (是) 或 'n' (否)")
            return ask_to_override(file_path)
    except KeyboardInterrupt:
        print("\n操作被用户中断")
        logger.info("用户中断了操作")
        sys.exit(0)
    except Exception as e:
        logger.error(f"询问覆盖时出错: {str(e)}", exc_info=True)
        return False

# -----------------------------
# 主入口
# -----------------------------
if __name__ == "__main__":
    start_time = time.time()
    logger.info("="*50)
    logger.info("身份证号脱敏工具启动")
    logger.info(f"启动时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
    logger.info("="*50)

    parser = argparse.ArgumentParser(description="身份证号脱敏工具（支持PDF和图片文件）")
    parser.add_argument("--input", required=True, help="输入路径（文件或文件夹）")
    parser.add_argument("--output", required=True, help="输出路径（文件或文件夹）")
    parser.add_argument("--dpi", type=int, help=f"转换分辨率（默认: {Config.PDF_DPI}）")
    parser.add_argument("--font", help=f"字体路径（默认: {Config.FONT_PATH}）")
    parser.add_argument("--font-scale", type=float, help=f"字体缩放比（默认: {Config.FONT_SCALE}，范围建议0.6-1.2）")
    args = parser.parse_args()

    # 更新配置
    if args.dpi: Config.PDF_DPI = args.dpi
    if args.font: Config.FONT_PATH = args.font
    if args.font_scale is not None: Config.FONT_SCALE = max(0.6, min(args.font_scale, 1.2))

    # 判断是文件还是文件夹
    if os.path.isfile(args.input):
        # 单文件处理
        if not os.path.isfile(args.input):
            logger.error(f"输入路径不是有效的文件: {args.input}")
            exit(1)
        os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True)
        process_file(args.input, args.output)
    elif os.path.isdir(args.input):
        # 文件夹处理
        if not os.path.isdir(args.input):
            logger.error(f"输入路径不是有效的文件夹: {args.input}")
            exit(1)
        os.makedirs(args.output, exist_ok=True)
        process_folder(args.input, args.output)
    else:
        logger.error(f"输入路径不存在: {args.input}")
        exit(1)

    end_time = time.time()
    logger.info("="*50)
    logger.info("身份证号脱敏工具运行结束")
    logger.info(f"结束时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
    logger.info(f"总运行时间: {end_time - start_time:.2f}秒")
    logger.info("="*50)